[TRTLLM-15160][feat] Add post-attention o_proj gate hook to shared MLA base and make trtllm::kda_decode inplace-only - #18375

Description

Two independent, low-risk PyTorch-backend refactors.

1. Make trtllm::kda_decode inplace-only — Drop the optional output tensor from the op schema so the CUDA kernel only ever writes into a caller-provided buffer.
   Allocation moves to the Python wrapper, and the KDA mixer's optimized decode path supplies a persistent, pool-slot-sized, never-reallocated output buffer
   (mirroring the existing _cs_dense CUDA-graph contract). The fake registration is updated in lockstep and the op's mutated tensors are registered in inplace_info
   (alias-declaration order). No behavior change today (KDA decode runs eagerly under CUDA graphs) — correct prep for torch.compile.

2. Add a post-attention o_proj gate hook to the shared MLA base — Introduce an overridable _apply_output_gate(hidden_states, attn_output) on the base MLA, applied
   to o_proj's input in MLA.forward before _project_output. Since o_proj is never fused into the register_to_config custom op, the hook runs identically for both
   the custom-op and eager branches, and before the helix-CP reduce-scatter. The base default is identity, so the DeepSeek/Eagle subclasses are unchanged. Kimi K3
   MLA overrides the hook with its sigmoid output gate and drops its forward() override, also picking up the base o_proj / helix-CP output-projection path.

Validation

• KDA decode unit tests (test_kda_decode.py, test_kda_decode_op.py): 77 passed on GB300 (SM103).
• Shared MLA base change is a strict no-op for the DeepSeek family — test_attention_mla.py gives byte-identical results with and without the change.
• Kimi K3 model-level parity: GSM8K DEP16/TP16 = 97.04, matching the ~96.8 reference within noise (no regression from the o_proj-gate refactor).

Notes

No API-signature changes and no new LLM args. The two changes are independent concerns and can be split into two PRs if preferred.

Summary

• Made trtllm::kda_decode strictly in-place.
• Moved output allocation to the Python wrapper.
• Added mutation metadata for compilation and updated fake registration.
• Added a persistent KDA decode output buffer for CUDA graph reuse.
• Added the _apply_output_gate hook to shared MLA.
• Moved Kimi K3 sigmoid gating into the hook.
• Preserved identity behavior for DeepSeek and Eagle MLA.
• Updated KDA decode and MLA hook tests.

Dev Engineer Review

• The C++, Python wrapper, fake implementation, and mutation metadata use consistent in-place semantics.
• The wrapper validates caller-provided output tensors and returns the populated buffer.
• The persistent KDA buffer supports optimized decode and CUDA graph pool sizing.
• The MLA hook runs before output projection and Helix reduce-scatter.
• The default hook preserves existing DeepSeek and Eagle behavior.
• No configuration files or test-list files changed.
• Reported parity results cover KDA decode, DeepSeek-family outputs, and Kimi K3 GSM8K.
• No additional correctness or regression issues are identified from the provided changes.

QA Engineer Review

• Modified tests:
  • tests/unittest/_torch/thop/parallel/test_kda_decode.py
  • tests/unittest/_torch/modules/test_mla_registry.py
• test_kda_decode now passes an explicit output buffer and validates the in-place result.
• MLA registry tests cover identity gating, custom gating order, and sparse MLA rejection.
• No corresponding test-db/ or qa/ coverage entries are identified in the provided changes.
• Verdict: needs follow-up.